package org.example;

import org.apache.flink.table.functions.ScalarFunction;
import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.JsonNode;
import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.ObjectMapper;

/**
 * 自定义UDF函数：从JSON中根据路径提取值
 * 支持类似$.data[0].id这样的路径格式，模拟JSONPath的核心功能
 * 
 * 此函数可以从嵌套的JSON结构中精确提取特定字段的值，支持:
 * - 对象字段访问: $.field
 * - 数组索引访问: $.array[0]
 * - 多级路径: $.object.field
 * - 数组中对象的字段: $.array[0].field
 * 
 * 使用场景：
 * 1. 从复杂的CDC消息中提取指定字段
 * 2. 访问深层嵌套的JSON结构
 * 3. 在不需要完整解析JSON为Map的情况下快速获取特定字段值
 * 
 * 示例：
 * 假设有JSON字符串:
 * ```
 * {
 *   "database": "ucenter",
 *   "table": "user_info",
 *   "type": "INSERT",
 *   "ts": 1623742201,
 *   "data": [
 *     {
 *       "id": "1001",
 *       "name": "张三",
 *       "age": 25
 *     }
 *   ]
 * }
 * ```
 * 
 * 使用示例及结果:
 * JSON_PATH_EXTRACT(json, '$.database') -> "ucenter"
 * JSON_PATH_EXTRACT(json, '$.data[0].id') -> "1001"
 * JSON_PATH_EXTRACT(json, '$.data[0].age') -> "25"
 * 
 * 在Flink SQL中的使用示例:
 * ```sql
 * -- 注册函数
 * CREATE TEMPORARY FUNCTION JSON_PATH_EXTRACT AS 'org.example.JsonPathExtractor';
 * 
 * -- 应用函数提取字段
 * SELECT 
 *   JSON_PATH_EXTRACT(json_column, '$.database') AS db_name,
 *   JSON_PATH_EXTRACT(json_column, '$.table') AS table_name,
 *   JSON_PATH_EXTRACT(json_column, '$.data[0].id') AS user_id,
 *   JSON_PATH_EXTRACT(json_column, '$.data[0].name') AS user_name
 * FROM source_table;
 * ```
 */
public class JsonPathExtractor extends ScalarFunction {
    
    private static final ObjectMapper objectMapper = new ObjectMapper();
    
    /**
     * 从JSON字符串中提取指定路径的值
     * 
     * 支持的路径格式：
     * - $.field         - 访问顶层字段
     * - $.nested.field  - 访问嵌套字段
     * - $.array[0]      - 访问数组元素
     * - $.array[0].field - 访问数组中对象的字段
     * 
     * 返回值为字符串，即使提取的是数字或布尔值，也会转换为字符串形式。
     * 如果路径不存在或解析错误，返回null。
     * 
     * @param jsonStr JSON字符串，例如: "{\"id\":\"1\",\"data\":[{\"name\":\"张三\"}]}"
     * @param path 路径表达式，例如: "$.data[0].name"
     * @return 提取的值字符串，如果路径无效则返回null
     */
    public String eval(String jsonStr, String path) {
        if (jsonStr == null || jsonStr.isEmpty() || path == null || path.isEmpty()) {
            return null;
        }
        
        try {
            // 解析JSON字符串
            JsonNode rootNode = objectMapper.readTree(jsonStr);
            
            // 移除路径开头的$符号
            if (path.startsWith("$")) {
                path = path.substring(1);
            }
            
            // 按.分割路径
            String[] pathParts = path.split("\\.");
            JsonNode currentNode = rootNode;
            
            // 逐级解析路径
            for (String part : pathParts) {
                if (currentNode == null) {
                    return null;
                }
                
                // 处理数组索引，如data[0]
                if (part.contains("[") && part.contains("]")) {
                    int startBracket = part.indexOf("[");
                    int endBracket = part.indexOf("]");
                    String fieldName = part.substring(0, startBracket);
                    int index = Integer.parseInt(part.substring(startBracket + 1, endBracket));
                    
                    // 先获取字段值
                    if (fieldName.length() > 0) {
                        currentNode = currentNode.get(fieldName);
                        if (currentNode == null) {
                            return null;
                        }
                    }
                    
                    // 再按索引获取数组元素
                    if (currentNode.isArray() && index >= 0 && index < currentNode.size()) {
                        currentNode = currentNode.get(index);
                    } else {
                        return null;
                    }
                } else {
                    // 普通字段
                    currentNode = currentNode.get(part);
                }
            }
            
            // 返回结果
            return currentNode != null ? currentNode.asText() : null;
        } catch (Exception e) {
            System.out.println("JSON路径提取异常: " + e.getMessage() + ", JSON: " + jsonStr + ", 路径: " + path);
            return null;
        }
    }
} 